package net.tp.struct.tree;

import java.util.*;
import java.util.function.*;

public class TreeTraversal {

    /**
     * Breadth first search.
     * @author Trung Phan
     * @param <N> node type
     * @param root the tree root
     * @param childrenGetter function to getchildren of a node
     * @param consumer consumer for each visited node
     */
    public static <N> void bfs(N root, Function<N, Iterable<N>> childrenGetter, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        if (root == null) return;

        Queue<N> queue = new LinkedList<>();
        queue.add(root);

        while (!queue.isEmpty()) {
            N u = queue.remove();
            consumer.accept(u);

            for (N w : childrenGetter.apply(u)) {
                if (w != null) queue.add(w);
            }
        }
    }

    /**
     * Depth first search. Very similar to breadth first search except that it's using stack instead of queue. The consumer must be placed where a node is popped, not when the node is pushed.
     *
     * Because the nature of the stack, the order in which nodes are visited is reverse. For the same depth first search traversal but in the correct order, see {@link TreeTraversal#preOrder}.
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root the tree root
     * @param childrenGetter function to getchildren of a node
     * @param consumer consumer for each visited node
     */
    public static <N> void dfs(N root, Function<N, Iterable<N>> childrenGetter, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        Stack<N> stack = new Stack<>();
        stack.push(root);

        while (!stack.empty()) {
            N u = stack.pop();
            consumer.accept(u);

            for (N w : childrenGetter.apply(u)) {
                if (w != null) stack.push(w);
            }
        }
    }


    /**
     * Pre-order Traversal. Very similar to depth first search. Care is needed to make sure the order of visiting children is not reverse. To make sure the order is not reverse, the children is temporary pushed into a temporary stack before being pushed to the main stack to reverse the order or pushing. For a pre-defined binary tree, we can simply push the right child node first before pushing the left child node.
     *
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root the tree root
     * @param childrenGetter function to getchildren of a node
     * @param consumer consumer for each visited node
     */
    public static <N> void preOrder(N root, Function<N, Iterable<N>> childrenGetter, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        Stack<N> tmp = new Stack<>(); // to read children in reverse order
        Stack<N> stack = new Stack<>();
        stack.push(root);

        while (!stack.empty()) {
            N u = stack.pop();
            consumer.accept(u);

            for (N w : childrenGetter.apply(u)) {
                if (w != null) tmp.push(w);
            }

            while (!tmp.empty()) stack.push(tmp.pop()); // reverse the order
        }
    }


    /**
     * Pre-order Traversal of Binary Tree.
     *
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root the tree root
     * @param consumer consumer for each visited node
     */
    public static <N extends AbstractBinNode<N>> void preOrder(N root, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        Stack<N> stack = new Stack<>();
        stack.push(root);

        while (!stack.empty()) {
            N u = stack.pop();
            consumer.accept(u);

            if (u.right != null) stack.push(u.right); // push right first
            if (u.left != null) stack.push(u.left);
        }
    }

    /**
     * In-order traversal.
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root tree root
     * @param consumer consumer for each visited node
     * @time $O(n)$
     * @space $O(h)$
     */
    public static <N extends AbstractBinNode<N>> void inOrder(N root, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        Stack<N> stack = new Stack<>();
        for (N u = root; u != null; u = u.left) stack.push(u); // push all left nodes to stack

        while (!stack.isEmpty()) {
            N u = stack.pop();
            consumer.accept(u);

            // push all left nodes of the right node
            for (u = u.right; u != null; u = u.left) stack.push(u);
        }
    }

    public static <N extends AbstractBinNode<N>> void eulerTour(N root, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        Stack<N> stack = new Stack<>();
        stack.push(root);
        N curr = root; // curr to determine if the traversal is up or down

        while (!stack.empty()) {
            N u = stack.peek();
            consumer.accept(u);

            if (u.isLeaf() || u.right == curr || (u.left == curr && u.right == null))
                stack.pop(); // either no child or both children are visited
            else if (u.left == null || u.left == curr)
                stack.push(u.right);
            else
                stack.push(u.left);

            curr = u;
        }
    }

    /**
     * Morris In-order traversal. This algoritm uses $O(1)$ space but runs in $O(n \log n)$ time. It modifies the current tree to create threaded link and restore the tree once the algorithm ends.
     *
     * <p>Recall threaded binary tree is the binary tree that utilizes the leaf nodes' left link to link to the successor's node for the in-order traversal.</p>
     *
     * <p>Morris' approach is to modify the tree to create those threaded link and remove them later. So the algorithm modifies the tree but restores it back once the algorithm completes.</p>
     *
     * <p>Note that if the tree has parent information, we don't need to use Morris' approach. It can be easy to traverse the tree using $O(1)$ space running in $O(n)$ time.</p>
     *
     *
     * <pre>
     *     5 ┄┄┄┄┄┄┄┄┐
     * ┌───┴───┐     ┊
     * 2       7     ┊
     *┌┴┐            ┊
     *  3            ┊
     * ┌┴┐           ┊
     *   4 ┄┄┄┄┄┄┄┄┄┄┘
     * </pre>
     *
     * <p>Maintaining the current node pointer, it's easy to traverse through node 2-3-4, but how can we jump from 4 to 5, that's when we need the threaded link from 4 to 5. Assuming the current node is 5, the predecessor node is the 4 (which is the right most node of node 5's the left chid). If this is the first time node 5 is visited, the link 4-5 does not exist, so we build the link 4-5 then set the current node to 2 (which is node 5's left node). Then the current node will traverse through 2-3-4-5. The second time node 5 is visited, the link 4-5 exists, we just need to remove this link and move the current node to the left (which is node 7).</p>
     *
     * <p>Algorithm:</p>
     * 
     * <p>The implementation below is for illustration only. Real implementation needs to take care of error handling to make sure the tree is restored if an exception is thrown</p>
     *
     * <pre>
     * Initialize current to root
     * While current not null
     *      If current has no left child
     *          Print current, move current to right child
     *      Else If current's predecessor node right link is null
     *          Set predecessor's right link to current
     *          Move current to left child
     *      Else
     *          Set predecessor's right link to null
     *          Print current, move current to right child
     * </pre>
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root tree root
     * @param consumer consumer for each visited node
     * @time $O(n \log n)$
     * @space $O(1)$
     */
    public static <N extends AbstractBinNode<N>> void morrisInOrder(N root, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        N curr = root;
        while (curr != null) {
            if (curr.left == null) {
                consumer.accept(curr);
                curr = curr.right;
            }
            else {
                N rightmost = curr.left; // right most of the left node
                while (rightmost.right != null && rightmost.right != curr)
                    rightmost = rightmost.right;

                if (rightmost.right == null) {
                    rightmost.right = curr; // add link
                    curr = curr.left;
                }
                else {
                    rightmost.right = null; // remove link
                    consumer.accept(curr);
                    curr = curr.right;
                }
            }
        }
    }

    /**
     * Morris Pre-order traversal. This algorithm is the modification of {@link #morrisInOrder} traversal. The principle is the same, the key difference is when to consume the node. This algoritm uses $O(1)$ space but runs in $O(n \log n)$ time. It modifies the current tree to create threaded link and restore the tree once the algorithm ends.
     *
     * <p>Note that if the tree has parent information, we don't need to use Morris' approach. It can be easy to traverse the tree using $O(1)$ space running in $O(n)$ time.</p>
     *
     * <p>The implementation below is for illustration only. Real implementation needs to take care of error handling to make sure the tree is restored if an exception is thrown</p>
     * 
     * <p>Algorithm:</p>
     *
     * <pre>
     * Initialize current to root
     * While current not null
     *      If current has no left child
     *          Print current, move current to right child
     *      Else If current's predecessor node right link is null
     *          Set predecessor's right link to current
     *          Print current, move current to left child
     *      Else
     *          Set predecessor's right link to null
     *          Move current to right child
     * </pre>
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root tree root
     * @param consumer consumer for each visited node
     * @time $O(n \log n)$
     * @space $O(1)$
     */
    public static <N extends AbstractBinNode<N>> void morrisPreOrder(N root, Consumer<N> consumer) {
        if (root == null || consumer == null) return;

        N curr = root;
        while (curr != null) {
            if (curr.left == null) {
                consumer.accept(curr);
                curr = curr.right;
            }
            else {
                N rightmost = curr.left; // right most of the left node
                while (rightmost.right != null && rightmost.right != curr)
                    rightmost = rightmost.right;

                if (rightmost.right == null) {
                    rightmost.right = curr; // add link
                    consumer.accept(curr);
                    curr = curr.left;
                }
                else {
                    rightmost.right = null; // remove link
                    curr = curr.right;
                }
            }
        }
    }

    /**
     * Post-order traversal.
     *
     * @author Trung Phan
     * @param <N> node type
     * @param root tree root
     * @param consumer consumer for each visited node
     */
    public static <N extends AbstractBinNode<N>> void postOrder(N root, Consumer<N> consumer) {
        if (root == null) return;
        
        Stack<N> stack = new Stack<N>();
        N curr = root;
        stack.push(curr);

        while (!stack.isEmpty()) {
            N u = stack.peek();
            if (u.isLeaf() || u.right == curr || u.left == curr) {
                curr = stack.pop(); // curr == u
                // the trick is this root here to know if we are traversing descending or ascending. For binary tree, there is only 2 scenario to check. If this is a generic tree, we have to iterate through all the children of this node. Otherwise, if we maintain the parent, we can find out the information in $O(1)$ time with u == root.parent. But if we know the parent, we don't need to stack solution, we can just use an easier solution which request less space

                consumer.accept(u);
            }
            else {
                // push in reverse order
                if (u.right != null) stack.push(u.right);
                if (u.left != null) stack.push(u.left);
            }
        }
    }


    /**
     * Generic recursive traversal.
     *
     * @param <N> node type
     * @param root tree root
     * @param childrenGetter function to get children of a node
     * @param preOrderConsumer visit nodes through pre-order traversal
     * @param postOrderConsumer visit nodes through post-order traversal
     * @param inOrderConsumer visit nodes through in-order traversal
     * @param eulerConsumer visit nodes following euler tour traversal
     */
    public static <N> void traverseRecursive(N root, Function<N, Iterable<N>> childrenGetter,
            Consumer<N> preOrderConsumer, Consumer <N> postOrderConsumer,
            BiConsumer <N, Integer> inOrderConsumer, Consumer<N> eulerConsumer
            ) {

        if (root == null) return;
        if (preOrderConsumer == null) preOrderConsumer = (n) -> {};
        if (postOrderConsumer == null) postOrderConsumer = (n) -> {};
        if (inOrderConsumer == null) inOrderConsumer = (n, i) -> {};
        if (eulerConsumer == null) eulerConsumer = (n) -> {};
        
        preOrderConsumer.accept(root);
        eulerConsumer.accept(root);

        int ithChild = -1; // needed for in order traversal
        for (N child : childrenGetter.apply(root)) {
            if (++ithChild > 0) inOrderConsumer.accept(root, ithChild-1);
            if (child != null) {
                traverseRecursive(child, childrenGetter, preOrderConsumer,
                        postOrderConsumer, inOrderConsumer, eulerConsumer);
                eulerConsumer.accept(root);
            }
        }

        if (ithChild == -1) inOrderConsumer.accept(root, 0);

        postOrderConsumer.accept(root);
    }


    /**
     * Generic traversal. This is the translation of the recursive version. It's not efficient, but it can do all the scenario that the recursive version do.
     *
     * @param <N> node type
     * @param root tree root
     * @param childrenGetter function to get children of a node
     * @param preOrderConsumer visit nodes through pre-order traversal
     * @param postOrderConsumer visit nodes through post-order traversal
     * @param inOrderConsumer visit nodes through in-order traversal
     * @param eulerConsumer visit nodes following euler tour traversal
     */
    public static <N> void traverse(N root, Function<N, Iterable<N>> childrenGetter,
            Consumer<N> preOrderConsumer, Consumer <N> postOrderConsumer,
            BiConsumer <N, Integer> inOrderConsumer, Consumer<N> eulerConsumer
            ) {

        if (root == null) return;
        if (preOrderConsumer == null) preOrderConsumer = (n) -> {};
        if (postOrderConsumer == null) postOrderConsumer = (n) -> {};
        if (inOrderConsumer == null) inOrderConsumer = (n, i) -> {};
        if (eulerConsumer == null) eulerConsumer = (n) -> {};

        class StackFrame {
            public N node;
            public Iterator<N> childIterator;
            public int ithChild; // needed for in order traversal
            public StackFrame(N node) {
                this.node = node;
                this.childIterator = childrenGetter.apply(node).iterator();
                ithChild = -1;
            }
        }

        Stack<StackFrame> stack = new Stack<>();

        stack.push(new StackFrame(root));
        preOrderConsumer.accept(root);

        boolean prevChildIsNull = false; // for euler traversal
        while (!stack.empty()) {
            StackFrame frame = stack.peek(); // only peek not pop
            N u = frame.node;
            if (!prevChildIsNull) eulerConsumer.accept(u);

            if (frame.childIterator.hasNext()) {
                N v = frame.childIterator.next();
                frame.ithChild++;
                if (frame.ithChild > 0) inOrderConsumer.accept(u, frame.ithChild-1);
                if (v != null) {
                    stack.push(new StackFrame(v));
                    preOrderConsumer.accept(v);
                }
                prevChildIsNull = (v == null);
            }
            else {
                if (frame.ithChild == -1) inOrderConsumer.accept(u, 0);
                stack.pop();
                postOrderConsumer.accept(u);
                prevChildIsNull = false;
            }
        }
        
    }

}
